﻿// my_reduce.cpp
#include <torch/extension.h>
using namespace at;  // Tensor 等都在 at 命名空间

Tensor my_reduce_fwd(const Tensor& y)
{
    // y 可以是任意 shape，这里举个最简单的：sum 规约
    return y.sum();  // 你可以替换成自己手写的 SIMD / OpenMP / CUDA 实现
}

// 反向：把上游梯度广播成 y 的形状
Tensor my_reduce_bwd(const Tensor& grad_out, const Tensor& y) { return grad_out.expand_as(y); }

// 注册正反两个符号
TORCH_LIBRARY(my_ops, m)
{
    m.def("my_reduce", my_reduce_fwd);
    m.def("my_reduce_backward", my_reduce_bwd);
}

// pybind11 让 Python 能 import
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
    m.def("forward", &my_reduce_fwd, "custom reduce forward");
    m.def("backward", &my_reduce_bwd, "custom reduce backward");
}